#include "jrm.h"

#include <assert.h>
/* */
#include <CL/cl.h>
#include <stdio.h>
#include <stdlib.h>

#define IS_INITIALIZED  (MA != NULL)
#define WAS_EXECUTED    (VX != NULL)

#define ABS(X)          ((X) > 0 ? (X) : -(X))

#define MIN(X, Y)       ((X) < (Y) ? (X) : (Y))

#define MAX(X, Y)       ((X) > (Y) ? (X) : (Y))

#define MAX_SOURCE_SIZE (0x100000)

static int ORDER    = -1;
static int ROW_TEST = -1;
static int ITE_MAX  = -1;
static int *ITER     = NULL;

static float ERROR  = 0;

static float *MA    = NULL;
static float *MA_S  = NULL;
static float *VB    = NULL;
static float *VB_S  = NULL;
static float *VX    = NULL;

static cl_mem D_MA_S = NULL;
static cl_mem D_VB_S = NULL;
static cl_mem D_VX = NULL;
static cl_mem D_ITER = NULL;
static cl_mem D_MITER = NULL;
static cl_mem D_ERROR = NULL;
static cl_mem D_ORDER = NULL;

/* Variables needed for OpenCL context load */
static cl_platform_id platform_id = NULL;
static cl_device_id device_id = NULL;   
static cl_uint ret_num_devices;
static cl_uint ret_num_platforms;
static cl_int ret;
static cl_context context;
static cl_command_queue command_queue;
static cl_program program;
static cl_kernel kernel;

void jrm_initialize(char *input)
{
    int i, j;

    /* Check initialization */
    if (IS_INITIALIZED)
        jrm_delete();

    ORDER = (int) strtol(input, &input, 10);
    ROW_TEST = (int) strtol(input, &input, 10);
    ERROR = (float) strtod(input, &input);
    ITE_MAX = (int) strtol(input, &input, 10);

    assert(ORDER > 0);

    /* Load GPU device */
    ret = clGetPlatformIDs(1,&platform_id,&ret_num_platforms);
    ret = clGetDeviceIDs(platform_id,CL_DEVICE_TYPE_GPU,1,&device_id,&ret_num_devices);
    context = clCreateContext(NULL,1,&device_id,NULL,NULL,&ret);
    command_queue = clCreateCommandQueue(context,device_id,0,&ret);

    /* Allocate memory for input matrix A and matrix A* */
    MA = (float *) malloc(ORDER * ORDER * sizeof(float));
    MA_S = (float *) malloc(ORDER * ORDER * sizeof(float));

    /* Allocate memory for input vector B and vector B* */
    VB = (float *) malloc(ORDER * sizeof(float));
    VB_S = (float *) malloc(ORDER * sizeof(float));

    /* Load data */
    for (i = 0; i < ORDER; i++)
    {
        for (j = 0; j < ORDER; j++)
        {
            MA[i * ORDER + j] = (float) strtod(input, &input);
        }
    }

    for (i = 0; i < ORDER; i++)
        VB[i] = (float) strtod(input, &input);
}

void jrm_delete()
{
    if (!IS_INITIALIZED)
        return;

    free(MA);
    free(MA_S);
    free(VB);
    free(VB_S);

    if (WAS_EXECUTED)
    {
        free(VX);
        free(ITER);

	/* Release OpenCL objects and memory */
	ret = clFlush(command_queue);
	ret = clFinish(command_queue);
	ret = clReleaseProgram(program);
	ret = clReleaseMemObject(D_VX);
	ret = clReleaseMemObject(D_MA_S);
	ret = clReleaseMemObject(D_VB_S);
	ret = clReleaseMemObject(D_ITER);
	ret = clReleaseMemObject(D_MITER);
	ret = clReleaseMemObject(D_ERROR);
	ret = clReleaseMemObject(D_ORDER);
	ret = clReleaseCommandQueue(command_queue);
	ret = clReleaseContext(context);
    }

    /* Set initial null values */
    ORDER    = -1;
    ROW_TEST = -1;
    ITE_MAX  = -1;
    ITER     = NULL;

    ERROR  = 0;

    MA = NULL;
    VB = NULL;
    MA_S = NULL;
    VB_S = NULL;
    VX  = NULL;
}

/* OpenCL kernel source as string */
const char *kernelSource =                                       "\n" \
"#define ABS(X)          ((X) > 0 ? (X) : -(X))                    \n" \
"__kernel void jrm_cuda_execute(  __global float *D_VX,            \n" \
"                       __global float *D_MA_S,                    \n" \
"                       __global float *D_VB_S,                    \n" \
"                       __global int *D_ITER,                      \n" \
"                       __global int *MAX_ITER,                    \n" \ 
"                       __global float *ERROR,                     \n" \
"                       __global int *ORDER)                       \n" \
"{                                                                 \n" \
"    __local float numerator, denominator;                         \n" \
"                                                                  \n" \
"    int i, j, iter;                                               \n" \
"    float x, a, b;                                                \n" \
"    i = get_global_id(0);                                         \n" \
"                                                                  \n" \
"    *ERROR = *ERROR * *ERROR;                                        \n" \
"                                                                  \n" \
"    if (i > *ORDER)                                                \n" \
"        return;                                                   \n" \
"                                                                  \n" \
"    b = D_VB_S[i];                                                \n" \
"    a = D_MA_S[i * *ORDER + i];                                    \n" \
"    for (iter = 0; iter < *MAX_ITER; iter++)                       \n" \
"    {                                                             \n" \
"        numerator = 0.0f;                                         \n" \
"        denominator = 0.0f;                                       \n" \
"                                                                  \n" \
"        /* Calculate new x value */                               \n" \
"        x = b;                                                    \n" \
"        for (j = 0; j < *ORDER; j++)                               \n" \
"            x -= D_MA_S[i * *ORDER + j] * D_VX[j];                 \n" \
"        x += a * x;                                               \n" \
"                                                                  \n" \
"        /* Update error */                                        \n" \
"        if (ABS(x - D_VX[i]) > numerator)                         \n" \
"            numerator = ABS(x - D_VX[i]);                         \n" \
"                                                                  \n" \
"        if (ABS(x) > denominator)                                 \n" \
"            denominator = ABS(x);                                 \n" \
"                                                                  \n" \
"        /* Update global x*/                                      \n" \
"        D_VX[i] = x;                                              \n" \
"                                                                  \n" \
"        barrier(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE);      \n" \
"                                                                  \n" \
"        if ((numerator / denominator) < *ERROR)                      \n" \
"            break;                                                \n" \
"    }                                                             \n" \
"                                                                  \n" \
"    D_ITER[i] = iter;                                             \n" \
"}                                                                 \n";

static void jrm_execute()
{
    int i, j;
    short finished = 0;

    VX = (float *) malloc(ORDER * sizeof(float));
    ITER = (int *) malloc(ORDER * sizeof(int));

    D_VX = clCreateBuffer(context, CL_MEM_READ_WRITE,ORDER * sizeof(float), NULL, &ret);
    D_MA_S = clCreateBuffer(context, CL_MEM_READ_ONLY,ORDER * ORDER * sizeof(float), NULL, &ret);
    D_VB_S = clCreateBuffer(context, CL_MEM_READ_ONLY,ORDER * sizeof(float), NULL, &ret);
    D_ITER = clCreateBuffer(context, CL_MEM_WRITE_ONLY,ORDER * sizeof(int), NULL, &ret);
    D_MITER = clCreateBuffer(context, CL_MEM_READ_ONLY,sizeof(int), NULL, &ret);
    D_ERROR = clCreateBuffer(context, CL_MEM_READ_ONLY,sizeof(float), NULL, &ret);
    D_ORDER = clCreateBuffer(context, CL_MEM_READ_ONLY,sizeof(int), NULL, &ret);

    for (i = 0; i < ORDER; i++)
        for (j = 0; j < ORDER; j++)
            if (i != j)
                MA_S[i * ORDER + j] = MA[i * ORDER + j] / MA[i * ORDER + i];
            else
                MA_S[i * ORDER + j] = 0;

    for (i = 0; i < ORDER; i++)
        VB_S[i] = VB[i] / MA[i * ORDER + i];

    ret = clEnqueueWriteBuffer(command_queue, D_VX, CL_TRUE, 0,ORDER * sizeof(float), VB_S, 0, NULL, NULL);
    ret = clEnqueueWriteBuffer(command_queue, D_MA_S, CL_TRUE, 0,ORDER * ORDER * sizeof(float), MA_S, 0, NULL, NULL);
    ret = clEnqueueWriteBuffer(command_queue, D_VB_S, CL_TRUE, 0,ORDER * sizeof(float), VB_S, 0, NULL, NULL);
    ret = clEnqueueWriteBuffer(command_queue, D_MITER, CL_TRUE, 0,sizeof(int), &ITE_MAX, 0, NULL, NULL);
    ret = clEnqueueWriteBuffer(command_queue, D_ERROR, CL_TRUE, 0,sizeof(float), &ERROR, 0, NULL, NULL);
    ret = clEnqueueWriteBuffer(command_queue, D_ORDER, CL_TRUE, 0,sizeof(int), &ORDER, 0, NULL, NULL);

    program = clCreateProgramWithSource(context, 1, (const char **)&kernelSource, NULL, &ret);
    ret=clBuildProgram(program,1,&device_id,NULL,NULL,NULL);

    if(ret!=CL_SUCCESS)
    {
	printf("Error: Failed to build program executable!\n");
	return;
    }

    size_t globalWorkSize = (size_t)ORDER;
    kernel = clCreateKernel(program, "jrm_cuda_execute", &ret);
    ret = clSetKernelArg(kernel,0,sizeof(cl_mem),(void *)&D_VX);
    ret = clSetKernelArg(kernel,1,sizeof(cl_mem),(void *)&D_MA_S);
    ret = clSetKernelArg(kernel,2,sizeof(cl_mem),(void *)&D_VB_S);
    ret = clSetKernelArg(kernel,3,sizeof(cl_mem),(void *)&D_ITER);
    //ret = clSetKernelArg(kernel,4,sizeof(cl_int),(void *)&ITE_MAX);
    ret = clSetKernelArg(kernel,4,sizeof(cl_mem),(void *)&D_MITER);
    ret = clSetKernelArg(kernel,5,sizeof(cl_mem),(void *)&D_ERROR);
    ret = clSetKernelArg(kernel,6,sizeof(cl_mem),(void *)&D_ORDER);

    ret = clEnqueueNDRangeKernel(command_queue, kernel, 1, NULL, &globalWorkSize, NULL, 0, NULL, NULL);
    ret = clEnqueueReadBuffer(command_queue, D_ITER, CL_TRUE, 0, ORDER * sizeof(int), ITER, 0, NULL, NULL);
    ret = clEnqueueReadBuffer(command_queue, D_VX, CL_TRUE, 0, ORDER * sizeof(float), VX, 0, NULL, NULL);
}

void jrm_evaluate(int *iter, int *row_test, float *estimated, float *real)
{
    int i;

    if (!IS_INITIALIZED)
        return;

    if (!WAS_EXECUTED)
        jrm_execute();

    *estimated = 0;
    *iter = 0;

    for (i = 0; i < ORDER; i++)
    {
        *estimated += MA[ROW_TEST * ORDER + i] * VX[i];
         if (*iter < ITER[i])
			*iter = ITER[i];
    }

    *real = VB[ROW_TEST];
    *row_test = ROW_TEST;
}
